from pymysql import Connection
from my_utils import Timer
import pandas as pd  # conda install pandas openpyxl


class Mysql(Timer):
    def __init__(self, **kwargs):
        super().__init__()
        self.database = kwargs.pop('mysql_db')
        self.db = Connection(
            user=kwargs.pop('mysql_user', 'root'),
            password=kwargs.pop('mysql_pwd'),
            host=kwargs.pop('mysql_host', 'localhost'),
            database=self.database,
            port=kwargs.pop('mysql_port', 3306),
            charset=kwargs.pop('mysql_charset', 'UTF8'),
        )
        self.cursor = self.db.cursor()

    def __del__(self):
        self.cursor.close()
        self.db.close()

    def commit(self, sql):
        self.debug(sql)
        try:
            self.cursor.execute(sql)
            self.db.commit()
        except Exception as e:
            print(e)

    def fetchall(self, query):
        self.cursor.execute(query)
        return self.cursor.fetchall()  # 有数据：tuple of tuple；无数据：()

    def fetchone(self, query, n=9999999):
        print(query)
        self.cursor.execute(query)
        for _ in range(n):
            one = self.cursor.fetchone()  # tuple
            if one:
                yield one

    def fetchone_dt(self, query, n=9999999):
        print(query)
        rowcount = self.cursor.execute(query)
        columns = [i[0] for i in self.cursor.description]
        length = len(columns)
        for _ in range(min(n, rowcount)):
            one = self.cursor.fetchone()  # tuple
            yield {columns[i]: one[i] for i in range(length)}

    def insert(self, dt, tb):
        ls = [(k, v) for k, v in dt.items() if v is not None]
        sql = 'INSERT %s (' % tb + ','.join(i[0] for i in ls) + \
              ') VALUES (' + ','.join('%r' % i[1] for i in ls) + ')'
        self.commit(sql)

    def update(self, dt, dt_condition, tb):
        sql = 'UPDATE %s SET ' % tb + ','.join('%s=%r' % (k, v) for k, v in dt.items()) \
              + ' WHERE ' + ' AND '.join('%s=%r' % (k, v) for k, v in dt_condition.items())
        self.commit(sql)

    def insert_or_update(self, dt, dt_condition, tb):
        """有则更新，无则插入"""
        query = 'SELECT * FROM ' + tb + ' WHERE ' + \
                ' AND '.join('%s=%r' % (k, v) for k, v in dt_condition.items())
        if self.fetchall(query) == ():
            dt.update(dt_condition)
            self.insert(dt, tb)
        else:
            self.update(dt, dt_condition, tb)

    def replace(self, dt, tb):
        """
        若存在相同数据，会先删除，再重新插入
        缺点：
            1、需要考虑全部数据列，不然会丢失部分列
            2、索引的原因，导致插入效率较低
        """
        ls = [(k, v) for k, v in dt.items() if v is not None]
        sql = 'REPLACE INTO %s (' % tb + ','.join(i[0] for i in ls) + \
              ') VALUES (' + ','.join('%r' % i[1] for i in ls) + ')'
        self.commit(sql)

    def truncate(self, tb):
        sql = 'TRUNCATE TABLE %s' % tb
        self.commit(sql)


class Excel(Mysql):
    def postfix(self, prefix):
        """导出Excel文件后缀名"""
        return prefix.replace('.xlsx', '') + '_{}.xlsx'.format(self.today)

    def db2sheet(self, query, prefix=''):
        """
        保存为单sheet的Excel文件
        :param query: 数据库查询语句
        :param prefix: 文件名前缀
        """
        df = pd.read_sql_query(query, self.db)
        if prefix:
            df.to_excel(self.postfix(prefix), index=False)
        return df

    def db2sheets(self, queries, prefix):
        """
        保存为多sheets的Excel文件
        :param queries: 字典{sheet名:sql}
        :param prefix: 文件名前缀
        """
        writer = pd.ExcelWriter(self.postfix(prefix))
        for sheet_name, query in queries.items():
            df = pd.read_sql_query(query, self.db)
            # if 'url' in df.columns:df['url'] = "'" + df['url']
            df.to_excel(writer, sheet_name=sheet_name, index=False)
        writer.save()

    def tb2sheet(self, tb):
        self.db2sheet('SELECT * FROM ' + tb, tb)
